// Copyright (c) 2020 cgsdfc
//
// This software is released under the MIT License.
// https://opensource.org/licenses/MIT

#pragma once

#include <string.h>  // memcpy
#include <type_traits>

#include "pyast/arena.h"
#include "pyast/ast_node.h"
#include "pyast/logging.h"

#include "asdl.h"  // Python header

#define PYAST_ENABLE_IF(...) \
  asdl_seq_internal::enable_if_t<__VA_ARGS__, int> = 0

// This file define the AsdlSeq<> template.
namespace pyast {
namespace asdl_seq_internal {

template <bool B, typename T = void>
using enable_if_t = typename std::enable_if<B, T>::type;

template <typename T>
using IsEnum = std::is_enum<T>;

template <typename T>
using IsAST = std::is_base_of<pyast::AstNode, T>;

template <typename T>
using IsObject = std::is_base_of<object, T>;

template <typename T>
using EnableIfEnum = enable_if_t<IsEnum<T>{}>;

template <typename T>
using EnableIfAST = enable_if_t<IsAST<T>{}>;

template <typename T>
using EnableIfObject = enable_if_t<IsObject<T>{}>;

using asdl_seq_ty = ::asdl_seq*;
using asdl_int_seq_ty = ::asdl_int_seq*;

// This RefWrapper requires T can be constructed from a raw ptr.
template <typename T, typename Seq, Seq (*CreateFunc)(Py_ssize_t, PyArena*)>
struct BasicAsdlSeqTraits {
  using underlying_type = Seq;
  using value_type = T;
  using reference = value_type&;
  using const_reference = const value_type&;
  using pointer = value_type*;
  using size_type = size_t;

  static size_type GetSize(underlying_type ptr) {
    PYAST_INTERNAL_DCHECK(ptr, "Null pointer");
    return asdl_seq_LEN(ptr);
  }
  static reference GetItemMutable(underlying_type ptr, size_type i) {
    PYAST_INTERNAL_DCHECK(i < GetSize(ptr), "Index out of range");
    auto index = static_cast<int>(i);
    return asdl_seq_GET(ptr, index);
  }
  static const_reference GetItem(underlying_type ptr, size_type i) {
    return const_cast<const_reference>(GetItemMutable(ptr, i));
  }
  static void SetItem(underlying_type ptr, size_type i, value_type val) {
    PYAST_INTERNAL_DCHECK(i < GetSize(ptr), "Index out of range");
    auto index = static_cast<int>(i);
    asdl_seq_SET(ptr, index, val);
  }
  static underlying_type NewSeq(Arena* arena, size_t size) {
    PYAST_INTERNAL_DCHECK(arena, "Null arena pointer");
    auto* result = CreateFunc(size, arena->ptr());
    if (!result)
      throw error_already_set();
    return result;
  }
  static underlying_type CopySeq(Arena* arena, underlying_type src) {
    auto size = GetSize(src);
    auto* dst = NewSeq(arena, size);
    ::memcpy(dst->elements, src->elements, sizeof(value_type) * size);
    return dst;
  }
};

struct VoidPointerSeqTraits
    : BasicAsdlSeqTraits<void*, asdl_seq_ty, ::_Py_asdl_seq_new> {};

struct IntSeqTraits
    : BasicAsdlSeqTraits<int, asdl_int_seq_ty, ::_Py_asdl_int_seq_new> {};

// This traits is capture information about the type to be held by AsdlSeq.
// It is not the type held by the C struct, but the C++ type we pretend to hold
// by using AsdlSeq<T>.
template <typename T, typename Enable = void>
struct AsdlSeqTraits {  // inherits from a basic version.
  // The type returns by `seq[i]` and `*it`, where `seq` is a AsdlSeq<T>, `it`
  // is a iterator.
  using reference = T&;

  // Like `reference`, but when the AsdlSeq is const-quelified.
  using const_reference = const T&;
};

// This traits handle subclasses of AST (AST nodes).
template <typename T>
struct AsdlSeqTraits<T, EnableIfAST<T>> : VoidPointerSeqTraits {
  template <bool Mutable>
  class ItemRef : public T {
    using value_type = typename T::underlying_type;

   public:
    using underlying_type = pointer;

    ItemRef(Arena* arena, underlying_type ptr)
        : T(arena, reinterpret_cast<value_type>(ptr)), ptr_(ptr) {}

    ItemRef& operator=(const T& rhs) {
      // PYAST_CHECK_ITEM_MUTABLE(Mutable);
      // PYAST_CHECK_CROSS_ARENA(!this->arena() || !rhs.arena() ||
      //                         this->arena() == rhs.arena());
      *ptr_ = rhs.ptr();
      return *this;
    }

   private:
    underlying_type ptr_;
  };

  using reference = ItemRef<true>;
  using const_reference = ItemRef<false>;
};

// This traits is for pybind11:object.
// TODO: test this specialization.
template <typename T>
struct AsdlSeqTraits<T, EnableIfObject<T>> : VoidPointerSeqTraits {
  template <bool Mutable>
  class ItemRef : public T {
   public:
    using underlying_type = pointer;

    ItemRef(Arena* arena, underlying_type ptr)
        : T(pybind11::reinterpret_borrow<T>(
              reinterpret_cast<::PyObject*>(ptr))),
          arena_(arena),
          ptr_(ptr) {}

    ItemRef& operator=(const T& obj) {
      // PYAST_CHECK_ITEM_MUTABLE(Mutable);
      *ptr_ = obj.ptr();
      arena_->AddObject(obj);
      return *this;
    }

   private:
    Arena* arena_;
    underlying_type ptr_;
  };

  using reference = ItemRef<true>;
  using const_reference = ItemRef<false>;
};

// This traits is for enum type or int.
template <typename T>
struct AsdlSeqTraits<T, EnableIfEnum<T>> : IntSeqTraits {
  template <bool Mutable>
  class ItemRef {
   public:
    using underlying_type = pointer;

    ItemRef(Arena, underlying_type ptr) : ptr_(ptr) {}

    operator T() const {
      // PYAST_INTERNAL_DCHECK(ptr_, "");
      return static_cast<T>(*ptr_);
    }

    ItemRef& operator=(const T& val) {
      // PYAST_CHECK_ITEM_MUTABLE(Mutable);
      *ptr_ = static_cast<int>(val);
      return *this;
    }

   private:
    underlying_type ptr_;
  };

  using reference = ItemRef<true>;
  using const_reference = ItemRef<false>;
};

}  // namespace asdl_seq_internal

// This is a wrapper of Python's AsdlSeq and asdl_int_seq.
template <typename T>
class AsdlSeq : public AstNode {
 public:
  using Traits = asdl_seq_internal::AsdlSeqTraits<T>;
  using value_type = T;
  using reference = typename Traits::reference;
  using const_reference = typename Traits::const_reference;
  using pointer = typename Traits::pointer;
  using size_type = typename Traits::size_type;
  using underlying_type = typename Traits::underlying_type;

  // Construct a null object.
  AsdlSeq() = default;

  AsdlSeq(const AsdlSeq&) = default;
  AsdlSeq& operator=(const AsdlSeq&) = default;

  AsdlSeq(underlying_type ptr, Arena* arena) : AstNode{ptr, arena} {}

  // Construct a seq and allocate zero'ed memory.
  AsdlSeq(size_type size, Arena* arena)
      : AsdlSeq{Traits::NewSeq(arena, size), arena} {}

  // Construct a seq from a range. Contents are copied into the new seq.
  template <typename Iter>
  AsdlSeq(Iter first, Iter last, Arena* arena)
      : AsdlSeq{std::distance(first, last), arena} {
    std::copy(first, last, begin());
  }

  // Construct a seq from initializer AsdlSeq.
  AsdlSeq(std::initializer_list<T> init, Arena* arena)
      : AsdlSeq{init.begin(), init.end(), arena} {}

  // Make AsdlSeq<Constant> => AsdlSeq<expr> possible.
  template <typename U, PYAST_ENABLE_IF(std::is_convertible<U, T>{})>
  AsdlSeq(const AsdlSeq<U>& other) : AsdlSeq{other.ptr(), other.GetArena()} {}

  // Access the underlying pointer.
  underlying_type ptr() const {
    return reinterpret_cast<underlying_type>(ptr_);
  }

  reference operator[](size_type i) {
    return reference(arena_, &Traits::GetItemMutable(ptr(), i));
  }

  const_reference operator[](size_type i) const {
    return const_reference(arena_, &Traits::GetItemMutable(ptr(), i));
  }

  reference at(size_type i) {
    if (i >= size())
      throw std::out_of_range("index out of range");
    return operator[](i);
  }

  const_reference at(size_type i) const {
    if (i >= size())
      throw std::out_of_range("index out of range");
    return operator[](i);
  }

  size_type size() const { return ptr() ? Traits::GetSize(ptr()) : 0; }
  bool empty() const { return size() == 0; }
  pointer data() const { return ptr() ? &ptr()->elements[0] : nullptr; }

 private:
  template <typename Ref>
  class Iterator : public std::iterator<std::forward_iterator_tag, T> {
   public:
    Iterator(Arena* arena, pointer ptr) : arena_(arena), ptr_(ptr) {}
    Ref operator*() const { return Ref(arena_, ptr_); }
    bool operator==(const Iterator& other) const { return ptr_ == other.ptr_; }
    bool operator!=(const Iterator& other) const { return !operator==(other); }
    Iterator& operator++() {
      ++ptr_;
      return *this;
    }

   private:
    Arena* arena_;
    pointer ptr_;
  };

 public:
  using iterator = Iterator<reference>;
  using const_iterator = Iterator<const_reference>;

  iterator begin() { return iterator{arena_, data()}; }
  iterator end() { return iterator{arena_, data() + size()}; }

  const_iterator begin() const { return const_iterator{arena_, data()}; }
  const_iterator end() const { return const_iterator{arena_, data() + size()}; }
};

}  // namespace pyast
